source('00_util_scripts/mod_seurat.R')
source('00_util_scripts/mod_bplot.R')

identify_indiv <- function(data, x){
  data |>
    filter(indiv == x) |>
    count(ind_cov, sort = TRUE) |>
    pull(ind_cov) |>
    pluck(1)
}

join_clues <- function(path, bid){
  freemx <- demux_sample(path)

overlp <- freemx |>
  mutate(.cell = str_remove(.cell, '-.+')) |>
  left_join(clues) |>
  filter(batch_id == bid)

nsample <- overlp$indiv |> unique() |> length()

overlp |>
  ggplot(aes(indiv, group = ind_cov, fill = ind_cov)) + geom_bar() +
  scale_fill_manual(values = DiscretePalette(20))

ind_id <- 0:(nsample-1) |> map_chr(\(x)identify_indiv(overlp,x))

maj <- tibble(indiv = 0:(nsample-1), majority = ind_id) |>
  right_join(overlp) |>
  filter(ind_cov == majority)

message('Fraction of dominant cluster: ',signif(nrow(maj) / nrow(overlp)))
message('Number valid barcodes: ',nrow(maj))
return(maj)}

clues <- read_csv('~/clues1.csv')

batch_read_id <- read_csv(
  'batch_cov,batch_id
  dmx_count_BH7YT2DMXX_YE_0907,1
  dmx_AbFlare-4,2
  dmx_YE_7-13,3')

clues <- clues |>
  select(-batch_id) |>
  left_join(batch_read_id)

clues |> count(batch_id)

clues |> filter(batch_id == 3) |>
  count(ind_cov_batch_cov)

# identify batch id --------
read_tsv('~/append-ssd/nextflowing/scrna-sle-perez2022-p32-34/cellranger/mtx_conversions/CLUES1_POOL03_2/barcodes.tsv', col_names = 'id') |>
  mutate(.cell = str_remove(id, '-.')) |>
  right_join(clues) |>
  filter(!is.na(id)) |>
  count(batch_cov, sort = TRUE)

# do all ind present in all lib in a pool?
dup_bc <- clues |>
  filter(batch_id == 2) |>
  count(.cell, ind_cov) |>
  filter(n > 1)

# do my freemuxlet result match with clues?
## pool1.2 ----------
freemx1.2 <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-3/cellranger/count/CLUES1_POOL01_2/outs/pool1.2.deconv.clust1.samples.gz',1)

## pool1.3 ----------
freemx1.3 <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-3/cellranger/count/CLUES1_POOL01_3/outs/pool1.3.deconv.clust1.samples.gz',1)

freemx1.3 |>
  ggplot(aes(indiv, group = ind_cov, fill = ind_cov)) + geom_bar() +
  scale_fill_manual(values = DiscretePalette(20))

# how freemuxlet umi affect result
freemx1.1 <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-re1/cellranger/count/CLUES1_POOL01_1/outs/pool1.1.deconv.clust1.samples.gz',1)

# how freemuxlet nsample affect result ----------
freemx2.1n19 <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-re1/cellranger/count/CLUES1_POOL02_1/outs/pool2.1n19.deconv.clust1.samples.gz', 2)

## if nsample=16 ?
freemx2.1n16 <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-re1/cellranger/count/CLUES1_POOL02_1/outs/pool2.1.deconv.clust1.samples.gz',2)

# how removing reads without UB or CB affect -------
freemx1.1 <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-re1/cellranger/count/CLUES1_POOL01_1/outs/pool1.1.deconv.clust1.samples.gz', 1)

freemx1.1cb <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-re1/cellranger/count/CLUES1_POOL01_1/outs/pool1.1n16.deconv.clust1.samples.gz', 1)

freemx1.1 |>
  ggplot(aes(indiv, group = ind_cov, fill = ind_cov)) + geom_bar() +
  scale_fill_manual(values = DiscretePalette(20))

g1<- last_plot()

freemx1.1cb |>
  ggplot(aes(indiv, group = ind_cov, fill = ind_cov)) + geom_bar() +
  scale_fill_manual(values = DiscretePalette(20))

g2<- last_plot()

g1 / g2

# cellranger vs starsolo? -------
freemx3.2cr <- join_clues('~/append-ssd/nextflowing/scrna-sle-perez2022-p32-34/cellranger/count/CLUES1_POOL03_2/outs/pool3.2n16.deconv.clust1.samples.gz', 3)

bc3.2ss <- read_tsv('~/append-ssd/nextflowing/scrna-sle-perez2022-sp32-34/star/CLUES1_POOL03_2/CLUES1_POOL03_2.Solo.out/Gene/filtered/barcodes.tsv.gz', col_names = '.cell')

clues |>
  filter(batch_id == 3 & .cell %in% bc3.2ss$.cell) |>
  filter(.cell %in% bc3.2cr) |>
  nrow()

bc3.2cr <- read_tsv('~/append-ssd/nextflowing/scrna-sle-perez2022-p32-34/cellranger/mtx_conversions/CLUES1_POOL03_2/barcodes.tsv', col_names = '.cell')

bc3.2cr <- bc3.2cr$.cell |> str_remove('-.')

clues |>
  filter(batch_id == 3 & .cell %in% bc3.2cr) |>
  nrow()

clues |>
  filter(batch_id == 3 & .cell %in% bc3.2ss$.cell & .cell %in% bc3.2cr) |>
  nrow()

# cellsnp-lite & vireosnp -------

#' Demultiplex metadata and expression matrix with vireosnp
#'
#' @param vireo_path Path to donor_ids.tsv of vireosnp outputs.
#' @param mex_path Path to directory containing respective scRNA-seq MEX files, for Read10X(). 
#' @param meta_data Data frame. Metadata containing `.cell`, `batch_cov`, `ind_cov`.
#' @param batch_id Character. Value in `meta_data$batch_cov` that match inputs. Leave blank to auto-detect.
#'
#' @returns List of:
#' 1. Filtered `meta_data` with added `donor_id` and modified `.cell` column.
#' 2. Filtered sparse matrix from `mex_path`
#' @export
#'
#' @examples
#' vireo_donor_id <- 'SRA14481836/vireo/donor_ids.tsv'
#' mex_path <- 'SRA14481836/star/Gene/filtered'
#' meta <- read_delim('DE_cells/results/perez_sle_sobj_meta.csv.gz')
#' # leave batch_id for auto-detection
#' meta_mex_list <- vireo_donor_id |>
#'   demux_sample(mex_path = mex_path, meta_data = meta)
#'   
demux_sample <- function(vireo_path, mex_path, meta_data, demux_n = 200, batch_id = NULL){
  # remove ambiguous & empty droplet
  vireo_id <- read_delim(vireo_path, col_types = 'cc------') |>
    filter(str_detect(donor_id, 'donor'))
  
  sra_id <- str_extract(vireo_path, 'SRA\\d+')
  
  # find all matched barcodes
  matched_barcodes <- meta_data |>
    mutate(cell = str_remove(.cell, '_.+'), .cell = NULL, orig.ident = NULL) |>
    inner_join(vireo_id)
  
  # auto-detect matched batch_id
  if (is.null(batch_id)) {
    batch_id <- matched_barcodes |>
      dplyr::count(batch_cov, sort = T) |>
      head(1) |>
      pull(batch_cov)
    
    message('Automaticly identified batch id: ', batch_id)
  }
  
  # filter for the only matched batch
  donor_count <- matched_barcodes |>
    filter(batch_cov == batch_id) |>
    mutate(n_cell_demux = n(), .by = c(donor_id, ind_cov))
  
  g1 <- donor_count |>
    distinct(donor_id, ind_cov, n_cell_demux) |>
    mutate(donor_id = fct_reorder(donor_id, n_cell_demux, .fun = max),
           ind_cov = fct_reorder(ind_cov, n_cell_demux, .fun = max)) |>
    ggplot(aes(donor_id, ind_cov, fill = n_cell_demux)) +
    geom_tile() +
    geom_text(aes(label = n_cell_demux), color = 'white') +
    labs(title = sra_id) +
    RotatedAxis()
  
  print(g1)
  
  # keep only matched donor(ind_cov)
  new_meta <- donor_count |>
    slice_max(n_cell_demux, by = ind_cov) |>
    filter(n_cell_demux > demux_n)
  
  # remove duplicated batch+ind+barcodes from multi-to-multi inner_join
  set_count <- new_meta |>
    dplyr::count(cell)
  
  duplicated_set <- set_count |>
    filter(n > 1) |>
    nrow()
  
  if(duplicated_set) {
    message(str_glue('Find {duplicated_set} duplicated barcode_ind_batch. They are removed.'))
    }
  
  unique_cell <- set_count |>
    filter(n == 1) |>
    pull(cell)
  
  filtered_mex <- Read10X(mex_path)
  
  valid_cell <- colnames(filtered_mex) |>
    intersect(unique_cell)
  
  message('Found barcodes with unique valid donor_id & quality: ',
          length(valid_cell))
  
  final_meta <- new_meta |>
    filter(cell %in% valid_cell) |>
    mutate(.cell = str_c(cell, ind_cov, batch_cov, sra_id, sep = '_'))
  
  message('Adding suffix to sparse matrix...')
  valid_set_name <- tibble(cell = valid_cell) |>
    left_join(final_meta) |>
    pull(.cell)
  
  final_mex <- filtered_mex[,valid_cell]
  
  colnames(final_mex) <- valid_set_name
  
  list(final_meta, final_mex)
}

meta <- read_delim('DE_cells/results/perez_sle_sobj_meta.csv.gz')

ena_meta <- read_tsv('mission/perez_2022_SLE/PRJNA728702.tsv')

downloaded_ena <- read_csv('~/append-ssd/freemuxlet/perez.starsolo.nf.csv')

## pool 836 ---------
vireo836 <- '~/append-ssd/freemuxlet/perez.836/SRA14481836/vireo/donor_ids.tsv'

mex836 <- '~/append-ssd/freemuxlet/perez.836/SRA14481836/Solo.out/GeneFull_Ex50pAS/filtered_full'

list836 <- vireo836 |>
  demux_sample(mex_path = mex836, meta_data = meta)

meta836 <- list836[[1]]

rna_counts <- list836[[2]] |> colSums2()

tibble(rna_counts = rna_counts) |>
  ggplot(aes(rna_counts)) +
  geom_density()

## SSD pool ---------------
vireo_paths <- list.dirs('~/append-ssd/freemuxlet') |>
  str_subset('vireo') |>
  map_chr(\(x)str_c(x,'/donor_ids.tsv'))

mex_paths <- list.dirs('~/append-ssd/freemuxlet') |>
  str_subset('filtered_full')

# 19 files only output ~2G data
ssd_vireo <- list(x = vireo_paths, y = mex_paths) |>
  pmap(\(x,y)demux_sample(vireo_path = x, mex_path = y, meta_data = meta),
       .progress = T)

ssd_meta <- ssd_vireo |>
  list_transpose() |>
  pluck(1) |>
  list_rbind()

ssd_meta |>
  ggplot(aes(nCount_RNA)) +
  geom_density()

rna_counts <- ssd_vireo |>
  list_transpose() |>
  pluck(2) |>
  map(colSums2)

tibble(rna_counts = list_c(rna_counts)) |>
  ggplot(aes(rna_counts)) +
  geom_density()

ssd_meta |>
  dplyr::count(.cell) |>
  filter(n == 1) |>
  nrow()

ssd_meta$ind_cov_batch_cov

## HDD pool ----------
vireo_paths <- list.dirs('~/mist/perez2022sle/') |>
  str_subset('vireo') |>
  map_chr(\(x)str_c(x,'/donor_ids.tsv'))

mex_paths <- list.dirs('~/mist/perez2022sle/') |>
  str_subset('filtered_full')

hdd_vireo <- list(x = vireo_paths, y = mex_paths) |>
  pmap(\(x,y)demux_sample(vireo_path = x, mex_path = y, meta_data = meta),
       .progress = T)

hdd_meta <- hdd_vireo |>
  list_transpose() |>
  pluck(1) |>
  list_rbind() |>
  bind_rows(ssd_meta)

hdd_meta |>
  dplyr::count(cell, batch_cov, ind_cov) |>
  ggplot(aes(batch_cov, fill = as.character(n) |> fct_inseq())) +
  geom_bar() +
  scale_fill_viridis_d()

plotly::ggplotly()

hdd_meta |>
  distinct(ind_cov_batch_cov, Status) |>
  dplyr::count(Status)

ourbatch <- hdd_meta$batch_cov |> unique()

meta |>
  filter(batch_cov %in% ourbatch)

hdd_meta |>
  distinct(n_cell_demux, ind_cov_batch_cov) |>
  ggplot(aes(n_cell_demux)) +
  geom_density()

meta |>
  mutate(cell = str_remove(.cell, '_.+')) |>
  distinct(batch_cov)

ssd_mex <- ssd_vireo |>
  list_transpose() |>
  pluck(2)

hdd_mex <- hdd_vireo |>
  list_transpose() |>
  pluck(2)

vireo_mex <- c(ssd_mex, hdd_mex)

vireo_bigmex <- vireo_mex[[1]] |>
  RowMergeSparseMatrices(vireo_mex[-1])

vireo_bigmex |> glimpse()

vireo_meta <- hdd_meta |>
  select(.cell, cell, batch_cov, ind_cov, ind_cov_batch_cov, Age, Sex, Status, SLE_status) |>
  column_to_rownames('.cell')

sobj <- vireo_bigmex |>
  CreateSeuratObject(meta.data = vireo_meta, min.cells = 3, min.features = 200)

sobj |>
  write_rds('mission/perez_2022_SLE/perez50sample.rds')

## batch2 ------------
barcode_pattern <- 'Solo.out/GeneFull_Ex50pAS/filtered_full/barcodes.tsv$'

barcode_batch2 <-
  list.files('~/append-ssd/freemuxlet/perez.batch2/', full.names = T,
             recursive = T, pattern = 'barcodes.tsv') |>
  str_subset(barcode_pattern)

barcode_batch3 <-
  list.files('~/mist/perez2022sle/perez_batch3/', full.names = T,
             recursive = T, pattern = 'barcodes.tsv') |>
  str_subset(barcode_pattern) 

barcode_batch4 <-
  list.files('~/mist/perez2022sle/perez_batch4/', full.names = T,
             recursive = T, pattern = 'barcodes.tsv') |>
  str_subset(barcode_pattern)

tidy_batch2 <- barcode_batch2 |>
  read_tsv(col_names = 'cell', id = 'batch') |>
  mutate(batch = str_extract(batch, 'SRA\\d+'))

tidy_batch3 <- barcode_batch3 |>
  read_tsv(col_names = 'cell', id = 'batch') |>
  mutate(batch = str_extract(batch, 'SRA\\d+'))

tidy_batch34 <- barcode_batch4 |>
  read_tsv(col_names = 'cell', id = 'batch') |>
  mutate(batch = str_extract(batch, 'SRA\\d+')) |>
  bind_rows(tidy_batch3)

meta_tidy <- meta |>
  mutate(cell = str_remove(.cell, '_.+'), batch_cov, ind_cov, .keep = 'none')

ena_meta <- read_tsv('mission/perez_2022_SLE/PRJNA728702.tsv')

downloaded_ena <- read_csv('~/append-ssd/freemuxlet/perez.starsolo.nf.csv')

ena_batch2 <- ena_meta |>
  mutate(batch = str_replace(run_accession, 'SRR', 'SRA'), sample_title, .keep = 'none') |>
  inner_join(tidy_batch2)

ena_batch34 <- ena_meta |>
  mutate(batch = str_replace(run_accession, 'SRR', 'SRA'), sample_title, .keep = 'none') |>
  inner_join(tidy_batch34)

batch34_pair <- ena_batch34 |>
  nest(data = -sample_title) |>
  pull(data, name = sample_title) |>
  map(\(x)left_join(x, meta_tidy, relationship = 'many-to-many') |>
        dplyr::count(batch_cov, sort = T) |> head(1)) |>
  list_rbind(names_to = 'sample_title')

batch_ind_count <- meta_tidy |>
  distinct(batch_cov, ind_cov) |>
  dplyr::count(batch_cov)

batch2_meta <- batch2_pair |>
  select(-n) |>
  inner_join(batch_ind_count) |>
  inner_join(ena_batch2) |>
  distinct(sample_title, n, batch, batch_cov) |>
  write_csv('mission/perez_2022_SLE/batch2_meta.csv')

batch34_meta <- batch34_pair |>
  select(-n) |>
  inner_join(batch_ind_count) |>
  inner_join(ena_batch34) |>
  distinct(sample_title, n, batch, batch_cov) |>
  write_csv('mission/perez_2022_SLE/batch34_meta.csv')

batch34_meta |>
  ggplot(aes(n, batch)) +
  geom_col()

b3_acc <- list.files('~/mist/perez2022sle/perez_batch3/', 'SRA')

batch34_meta |>
  filter(batch %in% b3_acc) |>
  mutate(sample = batch, vireo_n = n, .keep = 'none') |>
  write_csv('~/mist/perez2022sle/perez_batch3/input.csv')

batch34_meta |>
  filter(!(batch %in% b3_acc)) |>
  mutate(sample = batch, vireo_n = n, .keep = 'none') |>
  write_csv('~/mist/perez2022sle/perez_batch4/input.csv')

# merge into single rds ------
## pool3
demux_path <-
list.files(path = '~/append-ssd/nextflowing/',
           pattern = 'samples.gz',
           recursive = TRUE,
           full.names = TRUE)

pool123 <- tibble(
  path = demux_path,
  lib_name = demux_path |> str_extract('pool...'),
  pool_id = demux_path |> str_extract('(?<=pool).') |> as.numeric()
)

comp123 <- pool123

meta_comp123 <- map2(comp123$path, comp123$pool_id, join_clues) |>
  map2(comp123$lib_name, \(x,y)mutate(x, lib_name = y)) |>
  list_rbind()

meta_comp123 <- meta_comp123 |>
  mutate(.cell = str_c(lib_name, '_', .cell)) |>
  select(-...4)

meta_comp123 <- meta_comp123 |>
  count(.cell) |>
  filter(n == 1) |>
  left_join(meta_comp123)

## read h5 ---------
dir3_path <-
  list.files(path = '~/append-ssd/nextflowing/',
             pattern = 'filtered_feature_bc_matrix',
             include.dirs = TRUE,
             recursive = TRUE,
             full.names = TRUE) |>
  str_subset('h5', negate = TRUE)

h5.123 <- tibble(
  path = dir3_path,
  lib_name = demux_path |> str_extract('POOL....') |> str_to_lower() |>
    str_replace('_','.') |> str_remove('0'),
  pool_id = demux_path |> str_extract('(?<=POOL0).') |> as.numeric()
)

h5.123 <- h5.123 |>
  filter(str_detect(lib_name, '1.4|2.2'))

dir3_path <- set_names(h5.123$path, h5.123$lib_name)

# 4m46s to read pool1-3
system.time(merged_mtx <- Read10X(dir3_path, strip.suffix = TRUE))

glimpse(merged_mtx)

sobj123 <- CreateSeuratObject(merged_mtx, min.cells = 3, min.features = 200)

sobj_real <- sobj123 |> filter(.cell %in% meta_comp123$.cell)

sobj_real <- sobj_real |> left_join(meta_comp123)

sobj_real <- sobj_real |> select(-c(indiv, majority))

sobj_real |> write_rds('mission/perez-sle-pool1-3.rds')

sobj_real <- read_rds('mission/perez-sle-pool1-3.rds')

SeuratDisk::as.h5Seurat(sobj_real, 'mission/perez-sle-pool1-3.h5seurat')

sobj_real |>
  as.SingleCellExperiment() |>
  zellkonverter::writeH5AD('mission/perez-sle-pool1-3.h5ad')

sobj_real <- SeuratDisk::LoadH5Seurat('mission/perez-sle-pool1-3.h5seurat')

# finish remained stars ----------
batch2 <- read_csv('~/append-ssd/freemuxlet/starsolo.nf.b2.csv')

done_b2 <- list.files('~/append-ssd/freemuxlet/perez.batch2/', include.dirs = T)

batch2 |>
  filter(!(id %in% done_b2)) |>
  head(25) |>
  write_csv('~/mist/perez2022sle/starsolo.nf.b3.csv')

batch2 |>
  filter(!(id %in% done_b2)) |>
  tail(24) |>
  write_csv('~/mist/perez2022sle/starsolo.nf.b4.csv')
